"""
基于图神经网络的关联资产策略
策略特点：
1. 使用GNN建模资产间关联关系
2. 捕捉市场传染效应和联动性
3. 支持动态图结构学习
"""

import torch
import torch_geometric
from torch_geometric.nn import GCNConv
from torch_geometric.data import Data

class AssetGraphNetwork(torch.nn.Module):
    def __init__(self, num_features, hidden_dim=64):
        super().__init__()
        self.conv1 = GCNConv(num_features, hidden_dim)
        self.conv2 = GCNConv(hidden_dim, hidden_dim)
        self.predictor = torch.nn.Linear(hidden_dim, 1)
        
    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        
        # 图卷积层
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        
        # 全局池化
        x = torch_geometric.nn.global_mean_pool(x, data.batch)
        
        # 预测层
        return self.predictor(x)

class GraphTradingStrategy:
    def __init__(self, asset_names):
        self.asset_names = asset_names
        self.num_assets = len(asset_names)
        self.model = None
        self.correlation_threshold = 0.7
        
    def build_graph(self, returns_data):
        """构建资产关联图"""
        # 计算相关系数矩阵
        corr_matrix = np.corrcoef(returns_data.T)
        
        # 生成边列表
        edges = []
        for i in range(self.num_assets):
            for j in range(i+1, self.num_assets):
                if abs(corr_matrix[i,j]) > self.correlation_threshold:
                    edges.append([i,j])
                    edges.append([j,i])  # 无向图
                    
        edge_index = torch.tensor(edges, dtype=torch.long).t().contiguous()
        
        # 节点特征(标准化收益率和波动率)
        means = returns_data.mean(axis=0)
        stds = returns_data.std(axis=0)
        x = torch.tensor(np.column_stack([means, stds]), dtype=torch.float)
        
        return Data(x=x, edge_index=edge_index)
        
    def train(self, returns_data, epochs=100):
        """训练图神经网络"""
        graph_data = self.build_graph(returns_data)
        dataset = [graph_data] * len(returns_data)  # 简单示例
        
        self.model = AssetGraphNetwork(num_features=2)
        optimizer = torch.optim.Adam(self.model.parameters(), lr=0.01)
        criterion = torch.nn.MSELoss()
        
        for epoch in range(epochs):
            for data in dataset:
                optimizer.zero_grad()
                pred = self.model(data)
                loss = criterion(pred, torch.randn(1))  # 示例目标
                loss.backward()
                optimizer.step()
                
    def predict_portfolio(self, recent_returns):
        """预测最优组合权重"""
        if self.model is None:
            return np.ones(self.num_assets) / self.num_assets
            
        graph_data = self.build_graph(recent_returns)
        with torch.no_grad():
            scores = self.model(graph_data)
            weights = torch.softmax(scores, dim=0)
            return weights.numpy()

if __name__ == '__main__':
    # 示例用法
    assets = ['SPY', 'GLD', 'TLT', 'QQQ', 'IWM']
    strategy = GraphTradingStrategy(assets)
    
    # 模拟收益率数据
    np.random.seed(42)
    returns = np.random.randn(100, len(assets)) * 0.01
    
    # 训练模型
    strategy.train(returns)
    
    # 测试预测
    weights = strategy.predict_portfolio(returns[-30:])
    print("最优组合权重:", dict(zip(assets, weights)))